#! /usr/bin/env python

import os
import time
from group import GROUPS

DFF_ROOT = "/home/chengtao/dff7"
DFI_NAME = "fit"
INIT_PPF = "initial.ppf"
INPUT_PPF = "input.ppf"
OUTPUT_PPF = "output.ppf"
DFFFIT_EXE = "dfffit.exe"
EXCLUDE_FOLDER =[
#"01_Alkane", 
#"07_AldehydeKetone",
#"08_Acids", 
#"09_Amides",
]

class Parameter:
    def __init__(self, value = 0.0, fixed = False):
        self.value = value
        self.fixed = fixed

    def __str__(self):
        result = str(self.value)
        if (self.fixed):
            result += "*"
        return result

    def fix(self, fixed = True):
        self.fixed = fixed

def parseParameter(line):
    line = line.strip() # Remove white spaces
    fixed = False
    if line.endswith("*"):
        fixed = True
        line = line[:-1]
    return Parameter(float(line), fixed)

Parameter.parse = parseParameter

class ForceFieldItem:
    def __init__(self):
        self.term = "NONE"
        self.atoms = []
        self.pars = []
        self.tag = ""

    def __str__(self):
        result = self.term + ": " + ", ".join(self.atoms) + ": "
        for par in self.pars:
            result += str(par) + ", "
        result = result[:-2]
        if self.tag.strip() != "":
            result += ": " + self.tag
        return result

    def __eq__(self, item2):
        if self.term == item2.term:
            if len(self.atoms) == len(item2.atoms):
                for i in range(len(self.atoms)):
                    if self.atoms[i] != item2.atoms[i]:
                        return False
                return True
        return False

    def __ne__(self, item2):
        return not self.__eq__(item2)
                
    def fix(self, fixed = True):
        for i in self.pars:
            i.fix(fixed)

def parseForceFieldItem(line):
    line = line.strip()
    tokens = line.split(":", 3)
    result = ForceFieldItem()
    if len(tokens) > 0:
        result.term = tokens[0]
    if len(tokens) > 1:
        result.atoms = [i.strip() for i in tokens[1].split(",")]
    if len(tokens) > 2:
        result.pars = [parseParameter(i) for i in tokens[2].split(",")]
    if len(tokens) > 3:
        result.tag = tokens[3].strip()
    return result

ForceFieldItem.parse = parseForceFieldItem

class ForceField:
    def __init__(self):
        self.type = "NONE"
        self.data = []

    def fixAll(self):
        for i in self.data:
            i.fix(True)

    def write(self, file):
        file.write("#DFF:PPF\n#PROTOCOL = " + self.type+"\n")
        for item in self.data:
            file.write(str(item)+"\n")

    def index(self, item):
        for i in range(len(self.data)):
            if self.data[i] == item:
                return i
        return -1

    def __len__(self):
        return len(self.data)

def readForceField(file):
    result = ForceField()
    for line in file:
        line = line.strip()
        if line == "":
            break
        elif line.startswith("#DFF:PPF"):
            pass
        elif line.startswith("#PROTOCOL"):
            result.type = line.split('=')[1].strip()
        else:
            item = parseForceFieldItem(line)
            if item.term != "E0":
                result.data.append(item)
    return result
            
ForceField.read = readForceField

class EquivalenceTable:
    def __init__(self):
        self.data = { }

    def write(self, file):
        file.write("#AAT:\tNonbond\tBINC\tBond\tAngle/C\tAngle/S\tTorsion/C\tTorsion/S\tOOPA/C\tOOPA/S\n")
        for v in self.data.values():
            file.write(v[0] + ": " + "\t".join(v[1:]) + "\n")

    def fixAll(self):
        pass

    def __len__(self):
        return len(self.data)

def readEquivalenceTable(file):
    result = EquivalenceTable()
    for line in file:
        line = line.strip()
        if len(line) == 0:
            break
        elif line[0] != '#':
            tokens = line.split(": \t\n\r")
            tokens = [i.strip() for i in tokens]
            result.data[tokens[0]] = tokens
    return result
    
def fit(folder):
    print "####  Fitting", folder
    command = DFFFIT_EXE + ' "' + folder + '/' + DFI_NAME + '"'
    exe = os.popen(command, "r")
    for line in exe:
        print line.strip()
    exe.close()

def readLastOutputPPF(folder):
    ppfname = os.path.join(folder, OUTPUT_PPF)
    file = open(ppfname, "rt")
    ff = readForceField(file)
    file.close()
    return ff

def getEqtFileName(filename):
    index = filename.rfind(".")
    if index != -1:
        eqtname = filename[0:index]
    eqtname += ".eqt"
    return eqtname
    
def readLastOutputEQT(folder):
    eqtname = getEqtFileName(OUTPUT_PPF)
    eqtname = os.path.join(folder, eqtname)
    file = open(eqtname, "rt")
    eqt = readEquivalenceTable(file)
    file.close()
    return eqt

def readLastOutput(lastFolder):
    ff = readLastOutputPPF(lastFolder)
    eqt = readLastOutputEQT(lastFolder)
    return (ff, eqt)

def saveInput(folder, ff, eqt):
    ppfname = os.path.join(folder, INPUT_PPF)
    ppffile = open(ppfname, "wt")
    if ppffile:
        ff.write(ppffile)
        ppffile.close()

    eqtname = getEqtFileName(INPUT_PPF)
    eqtname = os.path.join(folder, eqtname)
    eqtfile = open(eqtname, "wt")
    if eqtfile:
        eqt.write(eqtfile)
        eqtfile.close()


def readInitPPF(folder):
    ppfname = os.path.join(folder, INIT_PPF)
    try:
        file = open(ppfname, "rt")
        ff = readForceField(file)
        file.close()
        return ff
    except IOError:
        return ForceField()
    
def readInitEQT(folder):
    eqtname = getEqtFileName(INIT_PPF)
    eqtname = os.path.join(folder, eqtname)
    print "$$$$$$$$", eqtname
    try:
        file = open(eqtname, "rt")
        print "$$$$$$$$", eqtname
        eqt = readEquivalenceTable(file)
        file.close()
        return eqt
    except IOError:
        return EquivalenceTable()

def readInitParameter(folder):
    ff = readInitPPF(folder)
    eqt = readInitEQT(folder)
    return (ff, eqt)

def mergeForceField(ff, ffI):
    if len(ffI) > 0:
        if ff.type != ffI.type:
            raise "Different Force Field Type"
        for i in ffI.data:
            if ff.index(i) != -1:
                raise "The item " + str(i) + " was fixed in input force field!"
            else:
                ff.data.append(i)
    return ff

def mergeEquivalenceTable(eqt, eqtI):
    for k, v in eqtI.data.iteritems():
        if k in eqt.data:
            raise "The EQT item " + k + " was fixed in input force field!"
        else:
            eqt.data[k] = v
    return eqt

def makeDFI(folder, template, fftype):
    global DFI_NAME, INIT_PPF, INPUT_PPF, OUTPUT_PPF
    DFI_NAME = "fit" + fftype
    INIT_PPF = "initial" + fftype + ".ppf"
    INPUT_PPF = "input" + fftype + ".ppf"
    OUTPUT_PPF = "output" + fftype + ".ppf"

    qmds = ""
    mols = ""

    for i in os.listdir("../Data/" + folder):
        if i.endswith(".qmd"):
            qmds += "  ../Data/" + folder + "/" + i + "\n"
            mols += "  ../Data/" + folder + "/" + i[:-4] + ".msd\n"
    qmds = qmds.strip()
    mols = mols.strip()

    if folder in EXCLUDE_FOLDER:
        print "The folder %s was excluded"%folder
        time.sleep(1) 
    else:
        template = template[:]
        template = template.replace("%DFF_ROOT%", DFF_ROOT)
        template = template.replace("%DATA_FILES%", qmds)
        template = template.replace("%MSD_FILES%", mols)
        template = template.replace("%FF_TYPE%", fftype)
        template = template.replace("%INPUT_FF%", folder + "/" + INPUT_PPF)
        template = template.replace("%OUTPUT_FF%", folder + "/" + OUTPUT_PPF)
        file = open(folder + "/" + DFI_NAME + ".dfi", "wt")
        file.write(template)
        file.close()

def copy_msd(folder, outfile):
    """This function is a temporary to copy the msd file generated by fit
progress. The purpose is to make input msd file for opt which has correct
atom type assignment. This function should be removed if the opt engine 
works correctly"""
    file_list = os.listdir("/home/chengtao/fitdff/FitAMBER/")
    for i in file_list:
        if i[-3:] == "msd":
            file_name = i[:-4]
            os.system("cp /home/chengtao/fitdff/FitAMBER/%s.msd /home/chengtao/fitdff/optData/%s/%s.msd"%(file_name, folder, file_name.split("_")[0]))
            outfile.write("cp /home/chengtao/fitdff/FitAMBER/%s.msd /home/chengtao/fitdff/optData/%s/%s.msd"%(file_name, folder, file_name.split("_")[0]))
    os.system("rm /home/chengtao/fitdff/FitAMBER/*.msd")
    
def autofit():
    fit(GROUPS[0])
    lastFolder = GROUPS[0]
    for folder in GROUPS[1:]:
        ff, eqt = readLastOutput(lastFolder)
        ff.fixAll()
        eqt.fixAll()

        ffI, eqtI = readInitParameter(folder)
        ff = mergeForceField(ff, ffI)
        eqt = mergeEquivalenceTable(eqt, eqtI)
        saveInput(folder, ff, eqt)
        fit(folder)
        lastFolder = folder

def autofit_copymsd():
    #add
    temp_file = open("/home/chengtao/fitdff/msdcopy.log", 'w')
    fit(GROUPS[0])
    #add
    copy_msd(GROUPS[0], temp_file)
    lastFolder = GROUPS[0]
    for folder in GROUPS[1:]:
        ff, eqt = readLastOutput(lastFolder)
        ff.fixAll()
        eqt.fixAll()

        ffI, eqtI = readInitParameter(folder)
        ff = mergeForceField(ff, ffI)
        eqt = mergeEquivalenceTable(eqt, eqtI)
        saveInput(folder, ff, eqt)
        fit(folder)
        #add
        copy_msd(folder, temp_file)
        lastFolder = folder
    temp_file.close()

def autofitAll_copymsd(type):
    file = open("fit.dfi", "rt")
    template = file.read()
    file.close()
    if not os.path.isdir("Fit" + type):
        os.mkdir("Fit" + type)
    os.chdir("Fit" + type)
    for i in GROUPS:
        if not os.path.isdir(i):
            os.mkdir(i)
        makeDFI(i, template, type)
    print DFI_NAME
    autofit_copymsd()
    os.chdir("..")

def autofitAll(type):
    file = open("fit.dfi", "rt")
    template = file.read()
    file.close()
    if not os.path.isdir("Fit" + type):
        os.mkdir("Fit" + type)
    os.chdir("Fit" + type)
    for i in GROUPS:
        if not os.path.isdir(i):
            os.mkdir(i)
        makeDFI(i, template, type)
    print DFI_NAME        
    autofit()
    os.chdir("..")


if __name__ == '__main__':
    autofitAll("AMBER")
    #autofitAll_copymsd("AMBER")
